import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import transforms

from models.vgg import VGGNet
from models.loss import EfficientSiameseLoss
from datasets.rank_dataloader_tid2013 import create_dataloader

def train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs=10, device='cuda'):
    """
    训练模型
    
    参数:
        model: 模型
        train_loader: 训练数据加载器
        val_loader: 验证数据加载器
        criterion: 损失函数
        optimizer: 优化器
        num_epochs: 训练轮数
        device: 训练设备
    """
    model.to(device)
    
    for epoch in range(num_epochs):
        # 训练阶段
        model.train()
        train_loss = 0.0
        
        for batch_idx, batch in enumerate(train_loader):
            inputs = batch['data'].to(device)
            labels = batch['label'].to(device)
            
            # 梯度清零
            optimizer.zero_grad()
            
            # 前向传播
            outputs = model(inputs)
            
            # 计算损失
            loss = criterion(outputs, labels)
            
            # 反向传播和优化
            loss.backward()
            optimizer.step()
            
            train_loss += loss.item()
            
            if batch_idx % 10 == 0:
                print(f'Epoch: {epoch+1}/{num_epochs}, Batch: {batch_idx}/{len(train_loader)}, Loss: {loss.item():.4f}')
        
        avg_train_loss = train_loss / len(train_loader)
        
        # 验证阶段
        model.eval()
        val_loss = 0.0
        
        with torch.no_grad():
            for batch in val_loader:
                inputs = batch['data'].to(device)
                labels = batch['label'].to(device)
                
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                
                val_loss += loss.item()
        
        avg_val_loss = val_loss / len(val_loader)
        
        print(f'Epoch: {epoch+1}/{num_epochs}, Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}')
        
        # 保存模型
        torch.save(model.state_dict(), f'checkpoint_epoch_{epoch+1}.pth')

def main():
    # 设置设备
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # 创建模型
    model = VGGNet(pretrained=True)
    
    # 定义损失函数 - 使用从Caffe转换的高效Siamese损失函数
    criterion = EfficientSiameseLoss(margin=10)
    
    # 定义优化器 - 使用与Caffe相似的学习率策略
    # 获取参数组，模拟Caffe中的lr_mult和decay_mult
    param_groups = model.get_parameter_groups()
    
    # 创建优化器 - 直接使用参数组中已设置的学习率和权重衰减
    optimizer = optim.SGD(param_groups, momentum=0.9)
    
    # 创建数据加载器
    train_loader = create_dataloader(
        pascal_root="/workspace/myRankIQA/data/score_tid2013",
        split="tid2013_train",
        batch_size=45,
        im_shape=(224, 224)
    )
    
    val_loader = create_dataloader(
        pascal_root="/workspace/myRankIQA/data/score_tid2013",
        split="tid2013_val",
        batch_size=45,
        im_shape=(224, 224)
    )
    
    # 训练模型
    train_model(
        model=model,
        train_loader=train_loader,
        val_loader=val_loader,
        criterion=criterion,
        optimizer=optimizer,
        num_epochs=10,
        device=device
    )

if __name__ == "__main__":
    main()